import argparse
import os
import random
import shutil
import time
import warnings
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.multiprocessing as mp
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from train import *
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
    

def plot_points(points, prednum,realname,size=0.1, axis=False):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    points=points.T
    # 绘制点云
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=size)
    
    if axis:
        ax.set_xlabel('X Label')
        ax.set_ylabel('Y Label')
        ax.set_zlabel('Z Label')
    plt.title(str(prednum)+" "+str(realname)+".png")
    plt.show()
    plt.savefig("test.png")
    
def print_args(args):
    print("\nArguments:")
    print("-" * 30)
    for key, value in vars(args).items():
        print(f"{key: <15}: {value}")
    print("-" * 30)

def parser_args():
    parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')

    parser.add_argument('--data_path', default='./data', type=str,
                        help='url used to set up distributed training')

    parser.add_argument('-n', '--num_class', default=10, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')
 
    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')
    parser.add_argument('-b', '--batch-size', default=64, type=int,
                        metavar='N',
                        help='mini-batch size (default: 256), this is the total '
                            'batch size of all GPUs on the current node when '
                            'using Data Parallel or Distributed Data Parallel')
    parser.add_argument('--weight', default='weight/best_model.pth', type=str,
                        help='url used to set up distributed training')
    parser.add_argument('--gpu', default=0, type=int,
                        help='GPU id to use.')
    parser.add_argument('-p', '--print-freq', default=10, type=int, metavar='N', 
                    help='print frequency (default: 10)')
    args = parser.parse_args()
    print_args(args)
    return args

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)

class ProgressMeter(object):
    def __init__(self, num_batches, meters, prefix=""):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = prefix

    def display(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        print('\t'.join(entries))

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'

def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res
    
def validate(val_loader, model, criterion, args):
    batch_time = AverageMeter('Time', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    progress = ProgressMeter(
        len(val_loader),
        [batch_time, losses, top1, top5],
        prefix='Test: ')

    # switch to evaluate mode
    model.eval()

    with torch.no_grad():
        end = time.time()
        for i, (points, target) in enumerate(val_loader):

            points,target =  points.cuda(args.gpu), target.cuda(args.gpu)
            points = points.transpose(2, 1)
            pred, _ = model(points)
            
            return 0
            pred_choice = pred.data.max(1)[1]
            
            loss = criterion(pred, target)

            # measure accuracy and record loss
            acc1, acc5 = accuracy(pred, target, topk=(1, 5))
            losses.update(loss.item(), points.size(0))
            top1.update(acc1[0], points.size(0))
            top5.update(acc5[0], points.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % args.print_freq == 0:
                progress.display(i)

        # TODO: this should also be done with the ProgressMeter
        print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'
              .format(top1=top1, top5=top5))
    # plot_points(points[0].cpu().numpy(),pred_choice[0].item(),target[0].item())    
    return top1.avg


def main(args):
    model = get_model(args.num_class)
    model.apply(inplace_relu)
    
    test_dataset = PointCloudDataset(root=args.data_path, split='test')
    val_loader = DataLoader(test_dataset, batch_size=1, shuffle=False,collate_fn=pad_collate_fn)

    torch.cuda.set_device(args.gpu)
    model = model.cuda(args.gpu)
    model.load_state_dict(torch.load(args.weight,map_location='cuda:0'))
    criterion = nn.CrossEntropyLoss().cuda(args.gpu)
    cudnn.benchmark = True
    top1_avg=validate(val_loader, model, criterion, args)
    print("=> top_1 AVG is  '{}'".format(top1_avg))

if __name__ == '__main__':
    args=parser_args()
    main(args)